#!/bin/bash

# 脚本名称：find_root_partition_offline.sh
# 功能说明：
# 本脚本用于查找离线磁盘（虚拟磁盘镜像或物理磁盘）上包含根文件系统的分区。
# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、partprobe、lsblk、df 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 识别磁盘上的所有分区，并尝试挂载它们。
# 5. 通过检测 /bin、/etc、/sbin 目录是否存在，判断分区是否为根文件系统。
# 6. 使用 trap 机制在脚本退出时自动卸载分区并断开 nbd 设备，确保系统资源得到及时释放。
#
# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、mount、umount。
# - 在执行操作前，请确保目标磁盘未被其他进程占用，以防数据损坏。

MOUNT_POINTS=()         # 存储所有挂载分区的挂载点
NBD_DEVICE="/dev/nbd0"
IS_VIRTUAL=false        # 标记是否为虚拟磁盘

# 安装依赖（适用于 Ubuntu/Debian）
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe" "lsblk" "df")
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    sudo apt-get install -y -qq qemu-utils parted >/dev/null

    for cmd in "${missing[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            echo "依赖安装失败，请手动安装：sudo apt-get install qemu-utils parted"
            exit 1
        fi
    done
}

# 检查磁盘镜像格式是否正确
check_disk_format() {
    local disk_path="$1"
    if ! qemu-img info "$disk_path" >/dev/null 2>&1; then
        echo "错误：无法识别的磁盘格式：$disk_path"
        exit 1
    fi
}

# 将虚拟磁盘镜像挂载到 NBD 设备上
setup_nbd() {
    local disk_path="$1"
    echo "加载 nbd 模块..."
    if ! sudo modprobe nbd max_part=16 2>/dev/null; then
        echo "nbd模块加载失败，尝试安装内核扩展模块..."
        sudo apt-get install -y "linux-modules-extra-$(uname -r)"
        if ! sudo modprobe nbd max_part=16; then
            echo "错误：无法加载 nbd 模块，请检查系统配置！"
            exit 1
        fi
    fi

    echo "连接 $disk_path 到 $NBD_DEVICE..."
    sudo qemu-nbd -c "$NBD_DEVICE" "$disk_path"
    sleep 2  # 等待设备初始化
}

# 挂载设备上的所有分区到 /mnt/ 下
mount_all_partitions() {
    local device="$1"
    echo "检测并挂载所有分区..."

    # 强制刷新分区表
    sudo partprobe "$device" >/dev/null 2>&1
    sleep 1

    # 获取所有分区（lsblk 输出分区名及类型）
    local partitions
    partitions=$(lsblk -rno NAME,TYPE "$device" | awk '$2=="part" {print $1}')

    if [[ -z "$partitions" ]]; then
        echo "未找到任何分区！"
        cleanup
        exit 1
    fi

    for part in $partitions; do
        local partition="/dev/$part"
        local mount_point="/mnt/$part"

        # 如果分区已经挂载，则获取其挂载点并加入检测列表
        mounted_point=$(findmnt -n -o TARGET "$partition")
        if [[ -n "$mounted_point" ]]; then
            echo "分区 $partition 已挂载到 $mounted_point，加入检测..."
            MOUNT_POINTS+=("$mounted_point")
            continue
        fi

        echo "挂载分区 $partition 到 $mount_point..."
        sudo mkdir -p "$mount_point"
        if ! sudo mount "$partition" "$mount_point"; then
            echo "挂载分区 $partition 失败，跳过该分区。"
            continue
        fi
        MOUNT_POINTS+=("$mount_point")
    done

    if [[ ${#MOUNT_POINTS[@]} -eq 0 ]]; then
        echo "错误：没有成功挂载任何分区。"
        cleanup
        exit 1
    fi

    echo "已挂载的分区："
    mount | grep "/mnt/"
}

# 判断挂载点是否包含根文件系统（简单检测 /bin、/etc、/sbin 是否存在）
is_root_partition() {
    local mount_point="$1"
    [[ -d "$mount_point/bin" && -d "$mount_point/etc" && -d "$mount_point/sbin" ]]
}

# 清理函数：卸载所有挂载的分区，并断开 NBD 连接（如果适用）
cleanup() {
    echo "开始清理..."
    for mount_point in "${MOUNT_POINTS[@]}"; do
        echo "卸载 $mount_point..."
        sudo umount "$mount_point" 2>/dev/null
        sudo rmdir "$mount_point" 2>/dev/null
    done

    if $IS_VIRTUAL; then
        echo "断开 qemu-nbd 连接..."
        sudo qemu-nbd -d "$NBD_DEVICE" >/dev/null 2>&1
    fi
}

# 捕获退出信号，执行清理
trap cleanup EXIT

# 主函数：解析参数、挂载分区并查找根文件系统所在分区
main() {
    if [[ $# -ne 1 ]]; then
        echo "用法: $0 <磁盘镜像文件或块设备>"
        echo "示例:"
        echo "  $0 ./disk.img   (虚拟磁盘镜像)"
        echo "  $0 /dev/sdb     (物理磁盘)"
        exit 1
    fi

    install_dependencies

    local disk_path="$1"
    local device

    if [[ -f "$disk_path" ]]; then
        IS_VIRTUAL=true
        check_disk_format "$disk_path"
        setup_nbd "$disk_path"
        device="$NBD_DEVICE"
    elif [[ -b "$disk_path" ]]; then
        if [[ $(lsblk -d -no TYPE "$disk_path") != "disk" ]]; then
            echo "错误：请指定整个磁盘设备（例如 /dev/sdb），而不是单个分区。"
            exit 1
        fi
        IS_VIRTUAL=false
        device="$disk_path"
    else
        echo "错误：无效的输入，必须是磁盘镜像文件或块设备。"
        exit 1
    fi

    mount_all_partitions "$device"

    local root_mount=""
    for mp in "${MOUNT_POINTS[@]}"; do
        if is_root_partition "$mp"; then
            root_mount="$mp"
            break
        fi
    done

    if [[ -n "$root_mount" ]]; then
        local partition_device
        partition_device=$(df "$root_mount" | tail -n 1 | awk '{print $1}')
        echo "----------------------------------------"
        echo "检测到根文件系统所在分区："
        echo "挂载点： $root_mount"
        echo "分区设备： $partition_device"
        echo "----------------------------------------"
        exit 0
    else
        echo "未找到包含完整根文件系统的分区。"
        exit 1
    fi
}

main "$@"

